package com.weave.tree;

import java.util.Comparator;

/**
 * @Description AVL树
 * @Author weave
 * @Date 2023/10/18 16:25
 */
public class AVLTree<E> extends BST<E> {
    public AVLTree(){
        this(null);
    }

    public AVLTree(Comparator<E> comparator){
        super(comparator);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null){
            if(isBalanced(node)){
                // 更新高度
                updateHeight(node);
            }else{
                rebalance2(node);
                // 整棵树回复平衡
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node){
        while((node = node.parent) != null){
            if(isBalanced(node)){
                // 更新高度
                updateHeight(node);
            }else{
                // 恢复平衡
                rebalance2(node);
            }
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    /**
     * 恢复平衡
     * @param grand
     */
    private void rebalance(Node<E> grand){
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        //是否是左结点
        if(parent.isLeftChild()){
            // LL
            if(node.isLeftChild()){
                rotate(grand,node,node.right,parent,parent.right,grand);
            }else {
                //LR
                rotate(grand,parent,node.left,node,node.right,grand);
            }
        }else{
            if(node.isLeftChild()){
                // RL
                rotate(grand,grand,node.left,node,node.right,parent);
            }else{
                // RR
                rotate(grand,grand,parent.left,parent,node.left,node);
            }
        }
    }

    /**
     * 旋转
     * @param r 谁失衡，谁是root
     * @param b
     * @param c
     * @param d
     * @param e
     * @param f
     */
    private void rotate(Node<E> r,Node<E> b,Node<E> c,Node<E> d,Node<E> e,Node<E> f){
        // 让d成为这可子树的根结点
        d.parent = r.parent;
        if(r.isLeftChild()){
            r.parent.left = d;
        }else if(r.isRightChild()){
            r.parent.right = d;
        }else{
            root = d;
        }

        // b-c
        b.right = c;
        if(c != null){
            c.parent = b;
        }
        updateHeight(b);

        // e-f
        f.left = e;
        if(e != null){
            e.parent = f;
        }
        updateHeight(f);

        // b-d-f
        d.left = b;
        d.right = f;
        b.parent = d;
        f.parent = d;
        updateHeight((d));
    }

    /**
     * 恢复平衡
     * @param grand 高度最低的那个不平衡结点
     */
    private void rebalance2(Node<E> grand){
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        if(parent.isLeftChild()){
            //LL
            if(node.isLeftChild()){
                // LL
                rotateRight(grand);
            }else{
                // LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        }else{
            // RL
            if(node.isLeftChild()){
                rotateRight(parent);
                rotateLeft(grand);
            }else{
                // RR
                rotateLeft(grand);
            }

        }
    }

    private void rotateLeft(Node<E> grand){
        Node<E> parent = grand.right;
        Node<E> child = parent.left;
        grand.right = child;
        parent.left = grand;
        afterRotate(grand,parent,child);
    }

    private void rotateRight(Node<E> grand){
        Node<E> parent = grand.left;
        Node<E> child = parent.right;
        grand.left = child;
        parent.right = grand;
        afterRotate(grand,parent,child);
    }

    private void afterRotate(Node<E> grand, Node<E> parent, Node<E> child){
        // 让parent称为子树的根节点
        parent.parent = grand.parent;
        if(grand.isLeftChild()){
            grand.parent.left = parent;
        }else if(grand.isRightChild()){
            grand.parent.right = parent;
        }else{
            // grand是root结点
            root = parent;
        }

        // 更新child的parent
        if(child != null){
            child.parent = grand;
        }

        // 更新grand的parent
        grand.parent = parent;

        // 更新高度
        updateHeight(grand);
        updateHeight(parent);


    }

    private boolean isBalanced(Node<E> node){
        return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }

    private void updateHeight(Node<E> node){
        ((AVLNode<E>)node).updateHeight();
    }

    /**
     * AVL树的节点，需要计算平衡因子，因此比普通二叉树多维护一个height属性(将height放入普通二叉树里没有用处，浪费空间)
     */
    private static class AVLNode<E> extends Node<E>{
        // height不可能是3，最开始是1，有一个左是2，再来左，就要旋转，然后是2；如果是右，则还是2，到不了3，因为先判断要不要旋转，再更新高度。
        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 获取该节点平衡因子(左子树高度 - 右子树高度)
         * @return int
         */
        public int balanceFactor(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        public Node<E> tallerChild(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            if(leftHeight > rightHeight){
                return left;
            }
            if(rightHeight > leftHeight){
                return right;
            }
            // 高度一样则返回同方向的，左子节点则返回左，否则返回右
            return isLeftChild()?left:right;
        }

        @Override
        public String toString() {
            String parentString = "null";
            if (parent != null) {
                parentString = parent.element.toString();
            }
            return element + "_p(" + parentString + ")_h(" + height + ")";
        }

    }
}
